// #pragma OPENCL EXTENSION cl_khr_fp16 : enable
// 不启用 cl_khr_subgroups

#define MAX_LOCAL 1024   // 每个 work-group 的最大线程数，可根据设备调整

__kernel void kernel_soft_max(
        __global char * src0,
        ulong offset0,
        __global char * src1,
        ulong offset1,
        __global char * src2,
        ulong offset2,
        __global char * dst,
        ulong offsetd,
        int ne00,
        ulong nb01,
        ulong nb02,
        ulong nb03,
        int ne12,
        int ne13,
        ulong nb11,
        ulong nb12,
        ulong nb13,
        ulong nb1,
        ulong nb2,
        ulong nb3,
        float scale,
        float max_bias,
        float m0,
        float m1,
        int n_head_log2
) {
    src0 = src0 + offset0;
    src1 = src1 + offset1;
    src2 = src2 + offset2;
    dst  = dst  + offsetd;

    int i03 = get_group_id(2);
    int i02 = get_group_id(1);
    int i01 = get_group_id(0);

    int i13 = i03 % ne13;
    int i12 = i02 % ne12;
    int i11 = i01;

    __global float * psrc0 = (__global float *)(src0 + i01*nb01 + i02*nb02 + i03*nb03);
    __global float * pmask = src1 != src0 ? (__global float *)(src1 + i11*nb11 + i12*nb12 + i13*nb13) : 0;
    __global float * psrc2 = src2 != src0 ? (__global float *)(src2) : 0;
    __global float * pdst  = (__global float *)(dst  + i01*nb1 + i02*nb2 + i03*nb3);

    float slope = 1.0f;

    // ALiBi 计算
    if (max_bias > 0.0f) {
        int h = i02;
        float base = h < n_head_log2 ? m0 : m1;
        int exp = h < n_head_log2 ? h + 1 : 2*(h - n_head_log2) + 1;
        slope = pow(base, exp);
    }

    const int lid = get_local_id(0);
    const int lsize = get_local_size(0);

    __local float scratch[MAX_LOCAL];

    // --- 第一阶段：局部最大值归约 ---
    float lmax = psrc2 ? psrc2[i02] : -INFINITY;
    for (int i00 = lid; i00 < ne00; i00 += lsize) {
        float v = psrc0[i00] * scale + (pmask ? slope * pmask[i00] : 0.0f);
        lmax = fmax(lmax, v);
    }

    scratch[lid] = lmax;
    barrier(CLK_LOCAL_MEM_FENCE);

    // 归约求最大值
    for (int stride = lsize >> 1; stride > 0; stride >>= 1) {
        if (lid < stride) {
            scratch[lid] = fmax(scratch[lid], scratch[lid + stride]);
        }
        barrier(CLK_LOCAL_MEM_FENCE);
    }

    float max_val = scratch[0];

    // --- 第二阶段：计算 exp 并求和 ---
    float lsum = 0.0f;
    for (int i00 = lid; i00 < ne00; i00 += lsize) {
        float val = psrc0[i00]*scale + (pmask ? slope*pmask[i00] : 0.0f);
        float e = exp(val - max_val);
        lsum += e;
        pdst[i00] = e;
    }

    scratch[lid] = lsum;
    barrier(CLK_LOCAL_MEM_FENCE);

    // 归约求和
    for (int stride = lsize >> 1; stride > 0; stride >>= 1) {
        if (lid < stride) {
            scratch[lid] += scratch[lid + stride];
        }
        barrier(CLK_LOCAL_MEM_FENCE);
    }

    float sum_val = scratch[0];

    // 若有 psrc2，lid==0 的线程将额外加上 exp(psrc2[i02] - max_val)
    if (psrc2 && lid == 0) {
        sum_val += exp(psrc2[i02] - max_val);
        scratch[0] = sum_val;
    }
    barrier(CLK_LOCAL_MEM_FENCE);

    sum_val = scratch[0];

    // --- 第三阶段：归一化 ---
    for (int i00 = lid; i00 < ne00; i00 += lsize) {
        pdst[i00] /= sum_val;
    }
}
